# -*- coding: utf-8 -*-
#### 人脸检测并识别计数 ####
# 依赖于ultralytics库
# 作者：Ekaros
# 版本：1.0




# 导入ultralytics库，可使用pip install ultralytics安装
from ultralytics import YOLO
import os

def get_faces(img_path, save_flag = False):
    model = YOLO('./model/face_detect.pt')
    result = model.predict(classes = 0, source = img_path, save = save_flag)
    return len(result[0].boxes.cls)

if __name__ == '__main__':
    path = './合照图片测试/'
    file_list = os.listdir(path)
    faces_list = []
    for file in file_list:
        img_path = path + file
        num = get_faces(img_path)
        faces_list.append(num)
        print(file, num)